"""
笑话评判系统 - 使用LangChain Expression Language实现 (LCEL风格)

这个脚本实现了一个笑话评判系统，它包含两个主要组件：
1. 笑话生成器(joker) - 根据给定主题生成笑话
2. 笑话评判员(critic) - 评价生成的笑话，指出其优缺点

该系统使用LangChain Expression Language的现代管道风格实现，基于Flowise工作流转换而来。
"""

from langchain.prompts import PromptTemplate
from langchain_ollama import ChatOllama
from langchain_core.runnables import RunnablePassthrough
from typing import Dict, Any

def create_joke_evaluation_system():
    """
    创建笑话评判系统的所有组件 - 使用LCEL管道风格
    
    返回:
        Dict: 包含系统所有组件的字典
    """
    # 初始化ChatOllama模型 - 用于生成笑话
    joke_model = ChatOllama(
        base_url="http://localhost:11434",
        model="qwen2.5:3b",
        temperature=0.7,
        streaming=True
    )

    # 初始化ChatOllama模型 - 用于评判笑话
    critic_model = ChatOllama(
        base_url="http://localhost:11434",
        model="qwen2.5:3b",
        temperature=0.7,
        streaming=True
    )

    # 定义生成笑话的Prompt模板
    joke_prompt = PromptTemplate(
        input_variables=["subject"],
        template="请根据下面的主题编写一段笑话，笑话的主题为{subject}"
    )

    # 定义评判笑话的Prompt模板
    critic_prompt = PromptTemplate(
        input_variables=["joke"],
        template="你是一名笑话大赛的评委，请根据下面的笑话，写出笑话的优缺点，笑话为{joke},将笑话的原文输出到评论的最前面"
    )

    # 创建生成笑话的链 - 使用LCEL管道风格
    joke_chain = joke_prompt | joke_model
    
    # 创建评判笑话的链 - 使用LCEL管道风格
    critic_chain = critic_prompt | critic_model
    
    # 创建完整的工作流 - 使用LCEL管道风格
    # 这个工作流将：
    # 1. 接收主题输入
    # 2. 生成笑话
    # 3. 将笑话传递给评判链
    # 4. 返回笑话和评价
    
    # 定义一个函数，将笑话生成结果转换为评判链需要的格式
    def format_for_critic(inputs):
        # 从joke_chain获取笑话
        joke = inputs["joke"]
        # 返回评判链需要的输入格式
        return {"joke": joke}
    
    # 完整工作流
    complete_workflow = (
        # 第一步：接收输入并生成笑话
        {"subject": lambda x: x["subject"]} 
        | joke_chain 
        | (lambda joke: {"joke": joke, "original_subject": joke})
        # 第二步：评判笑话
        | RunnablePassthrough.assign(
            critique=lambda x: critic_chain.invoke({"joke": x["joke"]})
        )
    )
    
    return {
        "joke_model": joke_model,
        "critic_model": critic_model,
        "joke_prompt": joke_prompt,
        "critic_prompt": critic_prompt,
        "joke_chain": joke_chain,
        "critic_chain": critic_chain,
        "complete_workflow": complete_workflow
    }

def joke_evaluation_workflow(subject: str, system_components: Dict[str, Any]) -> Dict[str, str]:
    """
    执行笑话评判工作流
    
    参数:
        subject: 笑话主题
        system_components: 系统组件字典
        
    返回:
        Dict: 包含笑话和评价的字典
    """
    # 方法1：使用单独的链
    # joke_chain = system_components["joke_chain"]
    # critic_chain = system_components["critic_chain"]
    
    # 生成笑话
    # joke = joke_chain.invoke({"subject": subject})
    # print(f"\n生成的笑话:\n{joke}\n")
    
    # 评判笑话
    # critique = critic_chain.invoke({"joke": joke})
    # print(f"笑话评价:\n{critique}")
    
    # return {
    #     "joke": joke,
    #     "critique": critique
    # }
    
    # 方法2：使用完整工作流（注释掉，可以取消注释使用）
    complete_workflow = system_components["complete_workflow"]
    result = complete_workflow.invoke({"subject": subject})
    print(f"\n生成的笑话:\n{result['joke']}\n")
    print(f"笑话评价:\n{result['critique']}")
    return result

def main():
    """主函数 - 运行笑话评判系统"""
    print("="*50)
    print("欢迎使用笑话评判系统！")
    print("该系统使用LangChain Expression Language的现代管道风格实现")
    print("基于Flowise工作流转换而来")
    print("="*50)
    print("请输入笑话主题，系统将生成笑话并进行评判。")
    print("输入 'quit' 或 'exit' 结束程序。")
    
    # 创建系统组件
    system_components = create_joke_evaluation_system()
    
    while True:
        user_input = input("\n请输入笑话主题: ")
        
        if user_input.lower() in ['quit', 'exit']:
            print("感谢使用笑话评判系统，再见！")
            break
            
        # 执行笑话评判工作流
        joke_evaluation_workflow(user_input, system_components)

if __name__ == "__main__":
    main()
